package com.agentframework.core.provider;

import com.agentframework.core.agent.AbstractAgent;
import com.agentframework.core.agent.AgentHolder;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.AfterReturning;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.DefaultToolExecutionResult;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * @author wangjun
 **/
@Aspect
public class ChatModelAspect {

    List<Message> toolHistoryMessages = Collections.synchronizedList(new ArrayList<>());

    @Before("execution(* org.springframework.ai.chat.model.ChatModel.stream(..))")
    public void beforeStream(JoinPoint joinPoint) {
        Prompt prompt = (Prompt) joinPoint.getArgs() [0];
        List<Message> instructions = prompt.getInstructions();
        instructions.add(0, instructions.remove(instructions.size() - 1));

        if (!toolHistoryMessages.isEmpty()) {
            List<Message> filterToolHistoryMessages = toolHistoryMessages.stream().filter(
                    tm -> (tm instanceof AssistantMessage && !((AssistantMessage) tm).getToolCalls().isEmpty())
                            || (tm instanceof ToolResponseMessage && !tm.getText().equals("Done"))).toList();
            prompt.getInstructions().addAll(filterToolHistoryMessages);
            toolHistoryMessages.clear();
        }
    }

    @AfterReturning(value = "execution(* org.springframework.ai.model.tool.ToolCallingManager.executeToolCalls(..))", returning = "result")
    public void afterExecuteToolCalls(JoinPoint joinPoint, Object result) {
        DefaultToolExecutionResult toolExecutionResult = (DefaultToolExecutionResult) result;
        List<Message> messages = toolExecutionResult.conversationHistory().stream().filter(item -> item.getMessageType() == MessageType.ASSISTANT || item.getMessageType() == MessageType.TOOL).toList();
        toolHistoryMessages.addAll(messages);
        ((AbstractAgent) AgentHolder.getAgent()).addToHistoryMessages(toolHistoryMessages);
    }

}
